package com.miku.ai.utils;

import com.miku.ai.common.ErrorCode;
import com.miku.ai.exception.BusinessException;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Slf4j
@Component
public class AiUtils {

    @Resource
    private ClientV4 client;

    private static final float STABLE_NUMBER = 0.05f;

    private static final float UN_STABLE_NUMBER = 0.95f;


    // 构造请求
//        List<ChatMessage> messages = new ArrayList<>();
//        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "作为一名营销专家，请为智谱开放平台创作一个吸引人的slogan");
//        messages.add(chatMessage);

    /**
     * 通用请求
     *
     * @param messages    消息列表
     * @param stream      是否流式返回
     * @param temperature 随机值(0-1)如果越高生成的答案越随机
     * @return
     */
    public String doRequest(List<ChatMessage> messages, Boolean stream, Float temperature) {
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(stream)
                .invokeMethod(Constants.invokeMethod)
                .temperature(temperature)
                .messages(messages)
                .build();
        // 调用
        ModelApiResponse invokeModelApiResp = null;
        try {
            invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        } catch (Exception e) {
            log.error("调用模型失败", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "调用模型失败");
        }
        ChatMessage result = invokeModelApiResp.getData().getChoices().get(0).getMessage();
        return result.getContent().toString();
    }

    /**
     * 通用请求-简化了消息传递
     *
     * @param sysMessage-系统消息
     * @param userMessage-用户消息
     * @param stream-          是否流式返回
     * @param temperature-     随机值(0-1)如果越高生成的答案越随机
     * @return
     */
    public String doRequest(String sysMessage, String userMessage, Boolean stream, Float temperature) {

        // 构造请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), sysMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doRequest(messages, stream, temperature);
    }

    /**
     * 同步消息
     *
     * @param sysMessage-系统消息
     * @param userMessage-用户消息
     * @param temperature-     随机值(0-1)如果越高生成的答案越随机
     * @return
     */
    public String doSyncRequest(String sysMessage, String userMessage, Float temperature) {
        return doRequest(sysMessage, userMessage, Boolean.FALSE, temperature);
    }

    /**
     * 同步消息-生成答案不稳定
     *
     * @param sysMessage-系统消息
     * @param userMessage-用户消息
     * @return
     */
    public String doSyncStableRequest(String sysMessage, String userMessage) {
        return doSyncRequest(sysMessage, userMessage, UN_STABLE_NUMBER);
    }

    /**
     * 同步消息-生成答案稳定
     *
     * @param sysMessage-系统消息
     * @param userMessage-用户消息
     * @return
     */
    public String doSyncUnStableRequest(String sysMessage, String userMessage) {
        return doSyncRequest(sysMessage, userMessage, STABLE_NUMBER);
    }

    /**
     * 通用请求-流式请求
     *
     * @param messages    消息列表
     * @param temperature 随机值(0-1)如果越高生成的答案越随机
     * @return
     */
    public Flowable<ModelData> doStreamRequest(List<ChatMessage> messages, Float temperature) {
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .invokeMethod(Constants.invokeMethod)
                .temperature(temperature)
                .messages(messages)
                .build();
        // 调用
        ModelApiResponse invokeModelApiResp = null;
        try {
            invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        } catch (Exception e) {
            log.error("调用模型失败", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "调用模型失败");
        }
        return invokeModelApiResp.getFlowable();
    }

    /**
     * 流式请求
     * @param sysMessage-系统消息
     * @param userMessage-用户消息
     * @param temperature-     随机值(0-1)如果越高生成的答案越随机
     * @return
     */
    public Flowable<ModelData> doStreamRequest(String sysMessage, String userMessage, Float temperature) {
        // 构造请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), sysMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doStreamRequest(messages, temperature);
    }


}
